#!/usr/bin/env python
# -*- encoding: utf-8 -*-
"""
@File    :   ADF_test.py
@Time    :   2021/12/01 14:55:51
@Author  :   glx 
@Version :   1.0
@Contact :   18095542g@connect.polyu.hk
@Desc    :   Do adf test
"""

# here put the import lib

from statsmodels.tsa.stattools import adfuller
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from copy import deepcopy


def ADF_test(ts_data):
    """
    Do adf test
    input: time series data: pd.Series
    output: adf test result and p-value
    """
    result = adfuller(ts_data)
    print("ADF Statistic: %f" % result[0])
    print("p-value: %f" % result[1])
    print("Critical Values:")
    for key, value in result[4].items():
        print("\t%s: %.3f" % (key, value))
    return result[1] < 0.05, result[1]


def do_ADF_test(smooth_data, step, Freq_range, ST):
    """
    Do adf test for all time series data
    input: smooth_data, step, Freq_range, ST
    ST < Freq_range
    output: adf test result and p-value

    test_result[key] = [bool]* len(smooth_data[key])
    """

    ts_data = {}
    test_result = {}
    # 对每个数据进行adf检验
    for key in smooth_data.keys():
        ts_data = smooth_data[key]
        result_list = []
        # 拆分为子段
        ts_data_list = []
        ts_data_length = len(ts_data)
        start = 0
        while start < ts_data_length:
            ts_data_list.append(ts_data[start : start + step])
            start += step

        for ts_sub_data in ts_data_list:
            # print(ts_sub_data)
            # 常数值处理
            if (ts_sub_data.std() < 0.001) or len(ts_sub_data) < 3:
                real_result = True
                result_list.append(
                    (real_result, ts_sub_data.index[0], ts_sub_data.index[-1])
                )
                continue
            # TODO check small number
            # 进行ADF检验
            adf_result = adfuller(ts_sub_data)
            # 存储子段的adf检验结果, 起点,终点
            # 检验策略:
            # 1. 如果p-value小于0.05, 则认为通过检验
            # 2. 如果极差大于Freq_range,则不通过
            # 3. 如果极差小于10, 则直接通过

            if (
                adf_result[1] < 0.05
                and max(ts_sub_data) - min(ts_sub_data) < Freq_range
            ):
                real_result = True
            elif max(ts_sub_data) - min(ts_sub_data) < ST:
                real_result = True
            else:
                real_result = False

            result_list.append(
                (real_result, ts_sub_data.index[0], ts_sub_data.index[-1])
            )

        test_result[key] = result_list
        stable_data = deepcopy(smooth_data)
        for key in test_result.keys():
            result_list = test_result[key]
            for result in result_list:
                if not result[0]:
                    stable_data[key].drop(
                        index=range(result[1], result[2] + 1), inplace=True
                    )

    return test_result, stable_data


if __name__ == "__main__":
    from smooth_data import do_smooth
    from visualize import visualize_single
    from get_matlab_data import get_data

    # 测试总数据
    data = get_data()
    smooth_data = do_smooth(data)
    test_result, stable_data = do_ADF_test(smooth_data, 50, 10)
    plt.figure(figsize=(20, 10))
    for key in test_result.keys():
        plt.plot(stable_data[key])

    # print(adf_result)
    # print(p_value)
    # plt.show()

    # ts_data = [1]*200 + np.random.randn(200)
    # ts_data[100:101] = 300 # 平稳序列
    # print("stationary:",ADF_test(ts_data)[0])
    # plt.figure()
    # plt.plot(ts_data)
